import sqlite3
from datetime import datetime

# 连接数据库
conn = sqlite3.connect('personal_finance.db')
cursor = conn.cursor()

# 创建收入和支出表
cursor.execute('''
CREATE TABLE IF NOT EXISTS Income (
    id INTEGER PRIMARY KEY,
    date TEXT,
    amount REAL,
    category TEXT,
    note TEXT)
''')
cursor.execute('''
CREATE TABLE IF NOT EXISTS Expense (
    id INTEGER PRIMARY KEY,
    date TEXT,
    amount REAL,
    category TEXT,
    note TEXT
)
''')
conn.commit()

# 记录收入
def add_income(date, amount, category, note):
    if amount <= 0:
        print("收入金额必须为正数。")
        return
    cursor.execute("INSERT INTO Income (date, amount, category, note) VALUES (?, ?, ?, ?)",
                   (date, amount, category, note))
    conn.commit()
    print("收入记录成功！")

# 记录支出
def add_expense(date, amount, category, note):
    if amount <= 0:
        print("支出金额必须为正数。")
        return
    cursor.execute("INSERT INTO Expense (date, amount, category, note) VALUES (?, ?, ?, ?)",
                   (date, amount, category, note))
    conn.commit()
    print("支出记录成功！")

# 查询收入或支出记录
def query_records(table, start_date=None, end_date=None, category=None):
    query = f"SELECT * FROM {table} WHERE 1=1"
    params = []
    if start_date:
        query += " AND date >= ?"
        params.append(start_date)
    if end_date:
        query += " AND date <= ?"
        params.append(end_date)
    if category:
        query += " AND category = ?"
        params.append(category)
    
    cursor.execute(query, params)
    records = cursor.fetchall()
    for record in records:
        print(record)

# 查看所有账单（包括收入和支出）
def view_all_records():
    print("\n--- 所有收入和支出记录 ---")
    
    print("\n收入记录：")
    cursor.execute("SELECT * FROM Income ORDER BY date")
    income_records = cursor.fetchall()
    for record in income_records:
        print(record)
    
    print("\n支出记录：")
    cursor.execute("SELECT * FROM Expense ORDER BY date")
    expense_records = cursor.fetchall()
    for record in expense_records:
        print(record)

# 主函数
def main():
    while True:
        print("\n--- 个人财务管理系统 ---")
        print("1. 记录收入")
        print("2. 记录支出")
        print("3. 查询收入记录")
        print("4. 查询支出记录")
        print("5. 查看所有账单")
        print("6. 退出")
        choice = input("请选择操作（1-6）：")
        
        if choice == '1':
            date = input("请输入日期（YYYY-MM-DD）：")
            amount = float(input("请输入金额："))
            category = input("请输入类别（如工资、奖金等）：")
            note = input("请输入备注：")
            add_income(date, amount, category, note)
        
        elif choice == '2':
            date = input("请输入日期（YYYY-MM-DD）：")
            amount = float(input("请输入金额："))
            category = input("请输入类别（如餐饮、交通、购物等）：")
            note = input("请输入备注：")
            add_expense(date, amount, category, note)
        
        elif choice == '3':
            print("查询收入记录")
            start_date = input("请输入开始日期（YYYY-MM-DD），留空则不限制：") or None
            end_date = input("请输入结束日期（YYYY-MM-DD），留空则不限制：") or None
            category = input("请输入类别，留空则不限制：") or None
            query_records("Income", start_date, end_date, category)
        
        elif choice == '4':
            print("查询支出记录")
            start_date = input("请输入开始日期（YYYY-MM-DD），留空则不限制：") or None
            end_date = input("请输入结束日期（YYYY-MM-DD），留空则不限制：") or None
            category = input("请输入类别，留空则不限制：") or None
            query_records("Expense", start_date, end_date, category)
        
        elif choice == '5':
            view_all_records()
        
        elif choice == '6':
            print("退出系统。")
            break
        
        else:
            print("无效选择，请重新选择。")

# 运行主程序
if __name__ == "__main__":
    main()

